import os
import torch


def get_files(directory, ext, with_dot=True):
    path_list = []
    for paths in [[os.path.join(dirpath, name) for name in filenames if
                   name.endswith(('.' if with_dot else '') + ext)] for
                  dirpath, dirnames, filenames in os.walk(directory)]:
        path_list.extend(paths)
    return path_list


class ParamsDataset(torch.utils.data.Dataset):
    def __init__(self, mode='train', **config):
        super(ParamsDataset, self).__init__()
        self.task_type = 'params'

        self.data = []
        if mode == 'train':
            with open(config['train_label_dir'] + '/params.txt', 'rt') as file:
                train_labels = [[float(text) for text in line.strip().split(',')] for line in file.readlines()]
            for label in train_labels:
                self.data.append(label)
            self.n_params = len(train_labels[0])
        else:
            with open(config['test_label_dir'] + '/params.txt', 'rt') as file:
                test_labels = [[float(text) for text in line.strip().split(',')] for line in file.readlines()]
            for label in test_labels:
                self.data.append(label)
            self.n_params = len(test_labels[0])
        

        self.config = config

    def __getitem__(self, index):
        params = self.data[index]
        return torch.Tensor(params),

    def __len__(self):
        return len(self.data)